#ifndef __KTRUSS_HPP__
#define __KTRUSS_HPP__

#include <unordered_map>
#include <utility>
#include <queue>
#include <vector>
#include <unordered_set>
#include <algorithm>
#include <iostream>
#include "graph.hpp"
using namespace std;
typedef vector<int> VI;
typedef vector<VI> VVI;
typedef pair<int, int> PI;
typedef vector<pair<unordered_map<int, VI>, unordered_map<int, int>>> TcpIndex;

struct uftree
{
    VI tr;
    uftree(int n) : tr(n)
    {
        for (int i = 0; i < n; i++)
            tr[i] = i;
    }
    int find(int x)
    {
        if (tr[x] == x)
            return x;
        return tr[x] = find(tr[x]);
    }
    bool uni(int x, int y)
    {
        int a = find(x), b = find(y);
        if (a == b)
            return false;
        tr[a] = b;
        return true;
    }
};

inline int ha(int a, int b)
{
    if (a > b)
        return (a << 16) | b;
    return (b << 16) | a;
}
inline PI deha(int h)
{
    return {h >> 16, h & 0xffff};
}

VVI get_edge_sup(const graph &g)
{
    int n = g.ns.size();
    VVI res = VVI(n, VI(n, 0));
    for (int i = 0; i < n; i++)
    {
        for (int j = i + 1; j < n; j++)
        {
            if (g.edgs[i][j] == "")
                continue;
            for (int k = 0; k < n; k++)
            {
                if (g.edgs[i][k] != "" && g.edgs[j][k] != "")
                    ++res[i][j], ++res[j][i];
            }
        }
    }
    return res;
}
VVI get_edge_truss(const graph &g)
{
    int n = g.ns.size();
    VVI res = VVI(n, VI(n, 0));
    VVI sup = get_edge_sup(g);
    unordered_map<int, int> um;
    for (int i = 0; i < n; i++)
    {
        for (int j = i + 1; j < n; j++)
        {
            if (g.edgs[i][j] != "")
            {
                um[ha(i, j)] = sup[i][j];
            }
        }
    }
    auto lam = [](const PI &a, const PI &b) { return a.second < b.second; };
    int k = 2;
    while (!um.empty())
    {
        for (; !um.empty();)
        {
            auto smal = min_element(um.begin(), um.end(), lam);
            if (smal->second > k - 2)
                break;
            auto xy = deha(smal->first);
            auto x = xy.first, y = xy.second;
            um.erase(smal);
            res[x][y] = res[y][x] = k;
            for (int i = 0; i < n; i++)
            {
                if (um.find(ha(i, x)) != um.end() && um.find(ha(i, y)) != um.end())
                {
                    --um[ha(i, x)];
                    --um[ha(i, y)];
                }
            }
        }
        ++k;
    }
    return res;
}
vector<VVI> query_simple_index(const graph &g, int k, string vq, const VVI &edgtruss)
{
    unordered_set<int> visited;
    vector<VVI> res;
    if (g.name_to_ind.find(vq) == g.name_to_ind.end())
    {
        cerr << "wrong vq" << endl;
        return res;
    }
    int target = g.name_to_ind.find(vq)->second;
    int n = g.ns.size();
    for (int i = 0; i < n; i++)
    {
        if (g.edgs[target][i] == "")
            continue;
        if (edgtruss[target][i] >= k && visited.find(ha(target, i)) == visited.end())
        {
            queue<int> q;
            q.push(ha(target, i));
            visited.insert(ha(target, i));
            VVI Cl(n, VI(n, 0));
            while (!q.empty())
            {
                auto xy = deha(q.front());
                q.pop();
                auto x = xy.first, y = xy.second;
                Cl[x][y] = Cl[y][x] = 1;
                for (int j = 0; j < n; j++)
                {
                    if (g.edgs[j][x] != "" && g.edgs[j][y] != "")
                    {
                        if (edgtruss[j][x] >= k && edgtruss[j][y] >= k)
                        {
                            auto tmp = ha(j, x);
                            if (visited.find(tmp) == visited.end())
                            {
                                visited.insert(tmp);
                                q.push(tmp);
                            }
                            tmp = ha(j, y);
                            if (visited.find(tmp) == visited.end())
                            {
                                visited.insert(tmp);
                                q.push(tmp);
                            }
                        }
                    }
                }
            }
            res.push_back(Cl);
        }
    }
    return res;
}
//返回值：VVI每行表示一个Tx，然后是节点编号和对应的儿子们，第二个unordered_map是储存边权重
TcpIndex get_tcp_index(
    const graph &g, const VVI &edgtruss)
{
    vector<pair<unordered_map<int, VI>, unordered_map<int, int>>> res;
    int n = g.ns.size();
    res.resize(n);
    for (int i = 0; i < n; i++)
    {
        unordered_map<int, int> Gx_w;
        //generate Gx and w
        for (int j = 0; j < n; j++)
        {
            if (j == i)
                continue;
            for (int m = 0; m < n; m++)
            {
                if (m == i)
                    continue;
                if (g.edgs[i][j] != "" && g.edgs[i][m] != "" && g.edgs[j][m] != "")
                {
                    Gx_w[ha(j, m)] = min(edgtruss[i][j], min(edgtruss[j][m], edgtruss[i][m]));
                }
            }
        }
        //init Tx
        auto &Ti = res[i];
        auto &Ti1 = Ti.first;
        auto &Ti2 = Ti.second;
        for (int j = 0; j < n; j++)
            if (g.edgs[i][j] != "")
                Ti1.insert(make_pair(j, VI()));
        int kmax = 0;
        for (auto &j : Gx_w)
            kmax = max(kmax, j.second);
        uftree uft(n);
        for (int k = kmax; k >= 2; k--)
        {
            for (auto &ele : Gx_w)
            {
                if (ele.second == k)
                {
                    auto xy = deha(ele.first);
                    auto x = xy.first, y = xy.second;
                    if (uft.uni(x, y))
                    {
                        Ti1[x].push_back(y);
                        Ti1[y].push_back(x);
                        Ti2[ele.first] = k;
                    }
                }
            }
        }
    }
    return res;
}

VI get_Vk(const TcpIndex &tcpind,int x,int y,int k)
{
    auto &x_index=tcpind[x];
    auto childmap=x_index.first;
    auto wmap=x_index.second;
    unordered_set<int> visited;
    queue<int> q;
    VI res;
    res.push_back(y);
    q.push(y);
    visited.insert(y);
    while(!q.empty()){
        auto a=q.front();q.pop();
        for(auto &b:childmap[a]){
            if(visited.find(b)==visited.end()&&wmap[ha(b,a)]>=k){
                res.push_back(b);
                visited.insert(b);
                q.push(b);
            }
        }
    }
    return res;
}

vector<VVI> query_tcp_index(const graph& g,int k,string vq,const VVI &edgtruss)
{
    int n=g.ns.size();
    VVI visited(n,VI(n,0));
    vector<VVI> res;
    if (g.name_to_ind.find(vq) == g.name_to_ind.end())
    {
        cerr << "wrong vq" << endl;
        return res;
    }
    int target=g.name_to_ind.find(vq)->second;

    TcpIndex tcpind=get_tcp_index(g,edgtruss);

    for(int i=0;i<n;i++)
    {
        if(g.edgs[i][target]!=""&&edgtruss[i][target]>=k&&visited[target][i]==0)
        {
            VVI Cl(n,VI(n,0));
            queue<PI> q;
            q.push({target,i});
            while(!q.empty()){
                auto ab=q.front();q.pop();
                auto a=ab.first,b=ab.second;
                if(visited[a][b]==0){
                    VI Vkab=get_Vk(tcpind,a,b,k);
                    for(auto &z:Vkab){
                        visited[a][z]=1;
                        Cl[a][z]=Cl[z][a]=1;
                        if(visited[z][a]==0)q.push({z,a});
                    }
                }
            }
            res.push_back(Cl);
        }
    }
    return res;
}

#endif
